"""请求跟踪中间件，生成和传递trace_id"""

import uuid
import logging
from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import Response

logger = logging.getLogger('middleware.trace')

class TraceMiddleware(BaseHTTPMiddleware):
    """请求跟踪中间件，为每个请求生成唯一的trace_id"""
    
    async def dispatch(self, request: Request, call_next) -> Response:
        # 从请求头中获取trace_id，如果不存在则生成一个新的
        trace_id = request.headers.get('X-Trace-Id', str(uuid.uuid4()))
        
        # 将trace_id添加到请求状态中，以便在请求处理过程中访问
        request.state.trace_id = trace_id
        
        try:
            # 执行请求处理
            response = await call_next(request)
            
            # 将trace_id添加到响应头中
            response.headers['X-Trace-Id'] = trace_id
            
            return response
        except Exception as e:
            # 记录异常信息，包含trace_id
            logger.error(f"请求处理异常 [trace_id={trace_id}]: {str(e)}")
            # 重新抛出异常，由异常处理中间件处理
            raise

def get_trace_id(request: Request = None) -> str:
    """从请求状态中获取trace_id，如果请求不存在或没有trace_id，则生成一个新的
    
    Args:
        request: FastAPI请求对象
        
    Returns:
        trace_id字符串
    """
    if request and hasattr(request.state, 'trace_id'):
        return request.state.trace_id
    
    # 如果没有请求对象或请求对象中没有trace_id，生成一个新的
    return str(uuid.uuid4())

__all__ = [
    'TraceMiddleware',
    'get_trace_id'
]